#!/usr/bin/env python
import os
import re
import os.path
import zlib
import types
import cPickle as myPickle

try:
    import sqlite3
except ImportError:
    sys.stderr.write('cannot load sqlite3 module.\n')

from threading import Lock
from Ganga.Utility.logging import getLogger

from AutoD3PDMakerTask import AutoD3PDMakerTask

class SQLiteTaskStore:
    """
    Persistent task store
    """
    def __init__(self, dbpath=''):

        self.dblock = Lock()

        self.dbpath = dbpath

        self.logger = getLogger('SQLiteTaskStore')

        if not self.dbpath:

            self.dbpath = os.path.join( os.getcwd(), '.d3pdmaker.db')

            if os.environ.has_key('AutoD3PDMakerTStore') and os.environ['AutoD3PDMakerTStore']:
                self.dbpath = os.environ['AutoD3PDMakerTStore']

        self.logger.info('Initializing SQLite task store connecton -> %s...' % self.dbpath)

        ## attempts to create database table
        self.createTables()

    def __pzdumps__(self, obj ):
        '''Pickle object, then compress the pickled.'''
        ## pickle_protocol=2, compress_level=5
        return zlib.compress( myPickle.dumps(obj, 2), 5)

    def __pzloads__(self, pzob ):
         '''Inverse of pzdumps:  decompress pz object, then unpickle.'''
         return myPickle.loads( zlib.decompress( pzob ) )

    def __eq__(self, other):

        if self.dbpath == other.dbpath:
            return True
        else:
            return False

    def connect(self):

        conn = None

        self.dblock.acquire()

        try:
            conn = sqlite3.connect(self.dbpath)
        except Exception, e:
            ## connection failed, release the dblock
            self.dblock.release()
            self.logger.error(e)
            raise e

        return conn

    def disconnect(self, conn):

        if not conn:
            pass
        else:
            try:
                conn.close()
            except Exception, e:
                self.logger.error(e)

            conn = None

        ## release the db lock
        self.dblock.release()

    def createTables(self):

        conn = None

        try:
            conn = self.connect()
            cur = conn.cursor()

            sqls = ["""CREATE TABLE IF NOT EXISTS tasks (id TEXT PRIMARY KEY, gtId INTEGER, input TEXT, output TEXT, status TEXT, lastUpdate TEXT, metadata TEXT, desc BLOB)"""]

            for sql in sqls:
                self.logger.debug(sql)
                cur.execute(sql)

        except Exception, e:
            self.logger.error(e)

        self.disconnect(conn)

    def updateTasks(self, tasks):
        """Adds new task entries into the database"""

        conn = None

        ick  = True

        try:

            conn = self.connect()

            cur  = conn.cursor()

            sql  = """INSERT OR REPLACE INTO tasks VALUES (?,?,?,?,?,?,?,?)"""

            for t in tasks:

                gtId = -1
                try:
                    gtId = t.gtId
                except AttributeError:
                    pass

                status = AutoD3PDMakerTask.status_map['new']
                try:
                    status = t.status
                except AttributeError:
                    pass

                md = {}
                try:
                    md = t.metadata
                except AttributeError:
                    pass

                output = {}
                try:
                    output = t.output
                except AttributeError:
                    pass

                id      = t.id
                input   = t.input
                ts      = t.lastUpdate
                desc    = sqlite3.Binary( self.__pzdumps__(t) )

                val = ( id, gtId, input, repr(output), status, ts, repr(md), desc )

                cur.execute( sql, val )

            conn.commit()

            ick = True

        except Exception, e:

            self.logger.error(e)

            if conn:
                conn.rollback()

            ick = False

        self.disconnect(conn)

        return ick

    def deleteTasks(self, clauses):
        """
        Deletes task entries matching the given clauses.
        """
        ## extract the metadata clauses to deal with it with a user SQL function
        mdclauses = {}
        if clauses.has_key('metadata'):
            mdclauses = clauses['metadata']
            del(clauses['metadata'])

        ## A user-defined SQL function for searching in metadata
        def mdMatch(mdstr):
            ## the metadata in task store is in string format
            md = eval(mdstr)
            for key,value in mdclauses.items():
                if not md.has_key( key ):
                    return 0
                elif md[key] != value:
                    return 0
            return 1

        conn = None

        ick = True

        try:
            conn = self.connect()
            if mdclauses:
                conn.create_function("mdMatch", 1, mdMatch)

            cur  = conn.cursor()

            sql  = 'DELETE FROM tasks WHERE 1=1 '

            if len(clauses) > 0:
                sql += ' AND ' + ' AND '.join( "%s=%s"%(key,repr(val)) for key,val in clauses.items() )

            if mdclauses:
                sql += ' AND mdMatch(metadata)=1 '

            self.logger.debug(sql)

            cur.execute( sql )

            conn.commit()

            ick = True

        except Exception, e:
            self.logger.error(e)
            ick = False

            if conn:
                conn.rollback()

        self.disconnect(conn)

        return ick

    def selectTasks(self, clauses):
        """
        Adds new task entries into the database. Keys of clauses dict. is the db table column name.

        Example:

           taskStore.selectTasks({'gtId'    : 1,
                                  'input'   : 'data09_900GeV.00140541.physics_MuonswBeam.merge.AOD.f170_m268',
                                  'metadata': {'version':1}})
        """

        ## extract the metadata clauses to deal with it with a user SQL function
        mdclauses = {}
        if clauses.has_key('metadata'):
            mdclauses = clauses['metadata']
            del(clauses['metadata'])

        ## A user-defined SQL function for searching in metadata
        def mdMatch(mdstr):
            ## the metadata in task store is in string format
            md = eval(mdstr)
            for key,value in mdclauses.items():
                if not md.has_key( key ):
                    return 0
                else:
                    if type( value ) in [ types.StringType ] and ( value.startswith('%') or value.endswith('%') ):
                        ## pattern matching for string kind variables
                        value_pat = re.sub(r'\%$', '', re.sub(r'^\%','',value))

                        ## always convert the md value to string type representation
                        value_md = repr(md[key])
                        
                        if value_md.find( value_pat ) < 0:
                            return 0
                    elif md[key] != value:
                        return 0
            return 1

        tasks = []
        conn  = None

        try:
            conn = self.connect()
            if mdclauses:
                conn.create_function("mdMatch", 1, mdMatch)

            cur  = conn.cursor()

            sql  = 'SELECT desc FROM tasks WHERE 1=1 '

            if len(clauses) > 0:
#                sql += ' AND ' + ' AND '.join( "%s=%s"%(key,repr(val)) for key,val in clauses.items() )

                for key,val in clauses.items():
                    op = '='

                    if type( val ) in [ types.StringType ] and ( val.startswith('%') or val.endswith('%') ) :
                        op = 'LIKE'

                    sql += ' AND %s %s %s' % (key,op,repr(val))

            if mdclauses:
                sql += ' AND mdMatch(metadata)=1 '

            sql += ' ORDER BY gtId'

            self.logger.debug(sql)

            for r in cur.execute( sql ):
                tasks.append( self.__pzloads__(r[0]) )
        except Exception, e:
            self.logger.error(e)
            tasks = []

        self.disconnect(conn)

        return tasks